import codecs
import random
from nltk.tokenize import sent_tokenize


random.seed(222)


prompt_path = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.format.prompt'
story_path = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.format.story'

# for seq2seq format, only story
result_src = '/home/rickwwang/project_own/story_generation/data/evaluate/coherence_src.txt'
result_tgt = '/home/rickwwang/project_own/story_generation/data/evaluate/coherence_tgt.txt'

# for lm format, prompt + story
result_lm = '/home/rickwwang/project_own/story_generation/data/evaluate/coherence_src_tgt.txt'

all_prompt = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_story = codecs.open(story_path, 'r', encoding='utf8').readlines()


def split_story(s):
    sens = sent_tokenize(s)
    return sens[:15]


srcs = []
tgts = []
lms = []

for p, s in zip(all_prompt, all_story):
    p = p.strip()
    s = s.strip()

    sens = split_story(s)
    if len(sens) < 15:
        continue
    srcs.append(p)
    tgts.append(s)
    lms.append(' <SEP> '.join([p, s]))

    for i in range(14):
        copy_sen = sens[:]
        copy_sen[i], copy_sen[i+1] = copy_sen[i+1], copy_sen[i]

        shuffle_s = ' '.join(copy_sen)

        srcs.append(p)
        tgts.append(shuffle_s)
        lms.append(' <SEP> '.join([p, shuffle_s]))


with codecs.open(result_src, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(srcs))
    fout.write('\n')
with codecs.open(result_tgt, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(tgts))
    fout.write('\n')
with codecs.open(result_lm, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(lms))
    fout.write('\n')
